# Copy from https://github.com/duongnphong/MaxPool2D-NumPy.git


import numpy as np
def add_padding(input_np, padding=0):
    if padding > 0:
        # Calculate the amount of padding needed on each side
        pad_amount = padding

        # Create an array for padded input_np
        padded_input_np = np.pad(
            input_np,
            ((0, 0), (pad_amount, pad_amount), (pad_amount, pad_amount)),
            mode="constant"
        )
    else:
        return input_np

    return padded_input_np


def maxpool(input_np, kernel_size: tuple, stride=None):
    stride = kernel_size[0] if stride is None else stride
    c, h, w = input_np.shape  # Height, width, and number of channels in the input_np
    output_channels = c  # Number of channels in the input_np
    output_height = (h - kernel_size[0]) // stride + 1
    output_width = (w - kernel_size[1]) // stride + 1

    output = np.zeros((output_channels, output_height, output_width))

    half_kernel_size = (kernel_size[0] // 2 , kernel_size[1] // 2)
    for channel in range(output_channels):
        for i in range(half_kernel_size[0], output_height + half_kernel_size[0]):
            for j in range(half_kernel_size[1], output_width + half_kernel_size[1]):
                # Extract subsection from channel 'channel' and (i*stride, j*stride) to ((i+1)*stride, (j+1)*stride)
                subsection = input_np[
                    channel,
                    i * stride -  half_kernel_size[0]: i * stride + half_kernel_size[0]+1,
                    j * stride -  half_kernel_size[1]: j * stride + half_kernel_size[1]+1,
                ]
                # Compute max value of the subsection and assign it to output
                output[channel, i-half_kernel_size[0], j-half_kernel_size[1]] = np.max(subsection)

    return output

def maxpool2d(
    input_np: np.ndarray,
    kernel_size,
    stride=None,
    padding=0,
) -> np.ndarray:
    """
    Perform a 2D MaxPool operation

    Args:
        input_np (np.ndarray): Input input_np
        kernel_size (int or tuple[int, int]): Size of the maxpooling window
        stride (int, optional): Stride value for the maxpooling operation. Default value is kernel_size.
        padding (int, optional): Padding value for the input input_np. Default is 0.

    Returns:
        np.ndarray: Resulting output of the maxpooling operation.


    Raises:
        TypeError: If `input_np` is not of type `numpy.ndarray`.
        ValueError: If `kernel_size` is invalid.
    """
    if not isinstance(input_np, np.ndarray):
        raise TypeError("input_np should be of type np.ndarray.")
    if isinstance(kernel_size, int):
        kernel_size = (kernel_size, kernel_size)
    elif isinstance(kernel_size, tuple) and len(kernel_size) == 1:
        kernel_size = (kernel_size[0], kernel_size[0])
    elif isinstance(kernel_size, tuple) and len(kernel_size) != 2:
        raise ValueError("Invalid kernel_size.")

    input_np = add_padding(input_np=input_np, padding=padding)
    # The `input_np` parameter in the provided code represents the input input_np on
    # which the maxpooling operation is performed. It is a NumPy array that
    # contains the pixel values of the input_np. The maxpooling operation involves
    # sliding a window (defined by `kernel_size`) over the input_np and taking the
    # maximum value within each window to downsample the input_np. The `maxpool2d`
    # function takes this input input_np, applies padding if specified, and then
    # performs the maxpooling operation using the `maxpool` function to generate
    # the output.
    out = maxpool(input_np=input_np, kernel_size=kernel_size, stride=stride)

    return out
